from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough

from langChain.config import model, embedding

#  文档
documents = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好而闻名",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="猫是独立的宠物，通常喜欢自己的空间",
        metadata={"source": "哺乳动物宠物文档"}
    ),
    Document(
        page_content="金鱼是初学者的流行宠物，只需要简单的护理",
        metadata={"source": "鱼类宠物文档"}
    ),
    Document(
        page_content="鹦鹉是聪明的宠物，能够模仿人类语言",
        metadata={"source": "鸟类宠物文档"}
    ),
    Document(
        page_content="兔子是社交宠物，需要足够的空间跳跃",
        metadata={"source": "哺乳动物宠物文档"}
    ),
]

#  实例化一个向量数据库，
vector_store = Chroma.from_documents(documents, embedding)


# # 本地化存储,persist_directory指定持久化目录
# vector_store = Chroma.from_documents(documents, embedding, persist_directory="./vector_store")
# # 加载磁盘中的向量数据库
# vectorstore = Chroma(persist_directory="./vector_store", embedding_function=embedding)

# # 相似度查询，并返回相似度分数，（分数低，相似度高）
# search_doc = vector_store.similarity_search_with_score("咖啡猫")
# print('检索器相似度查询:', search_doc)


# # 检索器,且只返回最高相似度的结构
# retriever = RunnableLambda(vector_store.similarity_search).bind(k=1)
# #  多个关键词的相似度搜索
# print('多个关键词的相似度搜索:', retriever.batch(["咖啡猫", "鲨鱼"]))


# 检索器,且只返回最高相似度的结构
retriever = RunnableLambda(vector_store.similarity_search).bind(k=1)

#  提示模板
message = '''
使用提供的上下文仅回答这个问题:
{question}
上下文:
{context}
'''
prompt_temp = ChatPromptTemplate.from_messages([('human', message)])
# RunnablePassthrough 允许将用户的问题之后再传递给prompt和model
chain = {'question': RunnablePassthrough(), 'context': retriever} | prompt_temp | model

resp = chain.invoke('请介绍一下猫？')

print(resp.content)
